"""
SSH连接管理模块

负责管理SSH连接池，提供连接复用、自动重连等功能
支持密码和密钥认证，提供统一的命令执行接口
"""

import paramiko
import logging
import os
import asyncio
from typing import Dict, Tuple, Optional
from threading import Lock
from .config_manager import Host,Task

# 获取当前模块的日志记录器
logger = logging.getLogger(__name__)


class SSHConnectionManager:
    """
    SSH连接池管理器
    
    功能：
    1. 管理多个SSH连接，实现连接复用
    2. 自动检测和重建断开的连接
    3. 提供线程安全的连接操作
    4. 统一的命令执行接口
    5. 连接状态监控和异常处理
    """
    
    def __init__(self):
        """
        初始化SSH连接管理器
        """
        # 连接池：存储活跃的SSH连接
        # key格式: "{host.name}|{host.host}:{host.port}|{host.username}"，value: paramiko.SSHClient对象
        self._connections: Dict[str, paramiko.SSHClient] = {}
        self._shells: Dict[str, paramiko.Channel] = {}
        
        # 线程锁：确保连接池操作的线程安全
        self._lock = Lock()
        
        # 连接统计
        self._connection_count = 0
        self._failed_connections = set()
    
    def get_connection(self, host: Host) -> paramiko.SSHClient:
        """
        获取SSH连接，支持连接复用和自动重连
        
        Args:
            host: 主机配置对象
            
        Returns:
            paramiko.SSHClient: SSH客户端对象
            
        Raises:
            Exception: 连接失败时抛出异常
        """
        # 生成连接键，用于标识唯一连接
        connection_key = f"{host.name}|{host.host}:{host.port}|{host.username}"
        try:
            
            # 使用线程锁确保连接池操作的原子性
            with self._lock:
                # 检查连接池中是否已存在连接
                if connection_key in self._connections:
                    client = self._connections[connection_key]
                    
                    # 验证现有连接是否仍然有效
                    if self._is_connection_alive(client):
                        logger.debug(f"复用现有连接: {connection_key}")
                        return client
                    else:
                        # 连接已断开，从池中移除
                        logger.warning(f"连接已断开，重新建立连接: {connection_key}")
                        self._remove_connection(connection_key)
                
                # 创建新连接
                client = self._create_connection(host)
                self._connections[connection_key] = client
                self._connection_count += 1
                
                logger.info(f"新建连接 [{self._connection_count}]: {connection_key}")
                return client
        except Exception as e:
            logger.error(f"连接失败: {connection_key}, 错误: {e}")
            raise
    
    def has_connection(self, key:str) -> bool:
        """
        检查是否存在指定的连接
        
        Args:
            key: 连接键（"{host.name}|{host.host}:{host.port}|{host.username}"）
            
        Returns:
            bool: 是否存在指定的连接
        """
        with self._lock:
            if key in self._connections:
                client = self._connections[key]
                if self._is_connection_alive(client):
                    return True
        return False

    def has_shell(self,key:str) -> bool:
        '''
        检查是否存在指定的shell连接
        Args:
            key: shell的key:{host.name}|{host.host}:{host.port}|{host.username}|{shell_inherit_task}
            
        Returns:
            bool: 是否存在指定的shell连接
        '''
        with self._lock:
            if key in self._shells:
                shell = self._shells[key]
                if shell and not shell.closed and shell.active:
                    return True
        return False

    def has_shell2(self, host: Host, key:str) -> bool:
        """
        检查是否存在指定的shell连接
        
        Args:
            host: 主机配置对象
            key: shell的id（任务名称）
            
        Returns:
            bool: 是否存在指定的shell连接
        """
        connection_key = f"{host.name}|{host.host}:{host.port}|{host.username}"
        shell_key = f"{connection_key}|{key}"
        with self._lock:
            if shell_key in self._shells:
                shell = self._shells[shell_key]
                if shell and not shell.closed and shell.active:
                    return True
        return False
    
    def get_shell(self, host: Host, key:str) -> paramiko.Channel:
        """
        获取SSH shell连接
        
        Args:
            host: 主机配置对象
            key: shell的id（任务名称）
            
        Returns:
            paramiko.Channel: SSH shell通道对象
        """
        connection_key = f"{host.name}|{host.host}:{host.port}|{host.username}"
        shell_key = f"{connection_key}|{key}"
        with self._lock:
            if shell_key in self._shells:
                shell = self._shells[shell_key]
                if shell and not shell.closed and shell.active:
                    return shell
        client = self.get_connection(host)
        shell = client.invoke_shell()
        self._shells[shell_key] = shell
        return shell
    
    def close_shell(self, host: Host, key:str):
        """
        关闭SSH shell连接
        
        Args:
            host: 主机配置对象
            key: shell的id（任务名称）
        """
        connection_key = f"{host.name}|{host.host}:{host.port}|{host.username}"
        shell_key = f"{connection_key}|{key}"
        with self._lock:
            if shell_key in self._shells:
                shell = self._shells[shell_key]
                shell.close()
                del self._shells[shell_key]
                logger.debug(f"关闭shell连接: {shell_key}")

    # 获取所有连接
    def get_all_connections(self) -> Dict[str, paramiko.SSHClient]:
        """
        获取所有活跃的SSH连接
        
        Returns:
            Dict[str, paramiko.SSHClient]: 所有活跃连接的字典，键为连接键，值为SSH客户端对象
        """
        with self._lock:
            return self._connections
    
    # 获取所有shell连接
    def get_all_shells(self) -> Dict[str, paramiko.Channel]:
        """
        获取所有活跃的SSH shell连接
        
        Returns:
            Dict[str, paramiko.Channel]: 所有活跃shell连接的字典，键为shell键，值为shell通道对象
        """
        with self._lock:
            return self._shells

    
    def _is_connection_alive(self, client: paramiko.SSHClient) -> bool:
        """
        检查SSH连接是否仍然有效
        
        Args:
            client: SSH客户端对象
            
        Returns:
            bool: 连接是否有效
        """
        try:
            # 获取传输层对象
            transport = client.get_transport()
            
            # 检查传输层是否存在且活跃
            if transport and transport.is_active():
                # 发送一个轻量级的请求测试连接
                transport.send_ignore()  # 发送ignore消息测试连接
                return True
            else:
                return False
                
        except Exception as e:
            logger.debug(f"连接检查失败: {str(e)}")
            return False


    def _create_connection(self, host: Host) -> paramiko.SSHClient:
        """
        创建新的SSH连接
        
        Args:
            host: 主机配置对象
            
        Returns:
            paramiko.SSHClient: SSH客户端对象
            
        Raises:
            Exception: 连接建立失败
        """
        # 创建SSH客户端实例
        client = paramiko.SSHClient()
        
        # 自动添加主机密钥（在生产环境中应谨慎使用）
        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        
        try:
            # 准备连接参数
            connect_params = {
                'hostname': host.host,
                'port': host.port,
                'username': host.username,
                'timeout': host.timeout,
                # 启用压缩以提高性能
                'compress': True,
                # 设置保活参数
                'sock': None,
            }
            
            # 根据配置选择认证方式
            if host.key_file:
                # 使用私钥认证
                if not os.path.exists(host.key_file):
                    raise FileNotFoundError(f"私钥文件不存在: {host.key_file}")
                connect_params['key_filename'] = host.key_file
                logger.debug(f"使用私钥认证: {host.key_file}")
                
            elif host.password:
                # 使用密码认证
                connect_params['password'] = host.password
                logger.debug("使用密码认证")
                
            else:
                raise ValueError(f"主机 {host.name} 未配置有效的认证方式")
            
            # 建立SSH连接
            client.connect(**connect_params)
            
            # 配置连接保活
            self._configure_keepalive(client)
            
            logger.info(f"成功连接到主机: {host.name}|{host.host}:{host.port}|{host.username}")
            
            # 从失败列表中移除（如果存在）
            connection_key = f"{host.name}|{host.host}:{host.port}|{host.username}"
            self._failed_connections.discard(connection_key)
            
            return client
            
        except paramiko.AuthenticationException as e:
            connection_key = f"{host.name}|{host.host}:{host.port}|{host.username}"
            error_msg = f"认证失败 [{connection_key}]: {str(e)}"
            logger.error(error_msg)
            self._failed_connections.add(connection_key)
            raise Exception(error_msg)
            
        except paramiko.SSHException as e:
            connection_key = f"{host.name}|{host.host}:{host.port}|{host.username}"
            error_msg = f"SSH连接失败 [{connection_key}]: {str(e)}"
            logger.error(error_msg)
            self._failed_connections.add(connection_key)
            raise Exception(error_msg)
            
        except Exception as e:
            connection_key = f"{host.name}|{host.host}:{host.port}|{host.username}"
            error_msg = f"连接主机失败 [{connection_key}]: {str(e)}"
            logger.error(error_msg)
            self._failed_connections.add(connection_key)
            # 清理失败的客户端
            try:
                client.close()
            except:
                pass
            raise Exception(error_msg)
    
    def _configure_keepalive(self, client: paramiko.SSHClient):
        """
        配置SSH连接的保活参数
        
        Args:
            client: SSH客户端对象
        """
        try:
            transport = client.get_transport()
            if transport:
                # 设置保活间隔（秒）
                transport.set_keepalive(30)
                logger.debug("已配置SSH保活参数")
        except Exception as e:
            logger.warning(f"配置保活参数失败: {str(e)}")
    
    # keep this function for demo
    def execute_command(self, host: Host, command: str, timeout: int = 30) -> Tuple[int, str, str]:
        """
        在指定主机上执行命令
        
        Args:
            host: 主机配置对象
            command: 要执行的命令
            timeout: 命令执行超时时间（秒）
            
        Returns:
            Tuple[int, str, str]: (退出状态码, 标准输出, 标准错误)
            
        Raises:
            Exception: 命令执行失败
        """
        # 获取SSH连接
        client = self.get_connection(host)
        
        try:
            logger.debug(f"执行命令 [{host.name}|{host.host}:{host.port}|{host.username}]: {command}")
            
            # 执行命令
            stdin, stdout, stderr = client.exec_command(
                command, 
                timeout=timeout,
                get_pty=False  # 不分配伪终端，避免交互问题
            )
            
            # 等待命令执行完成并获取退出状态码
            exit_status = stdout.channel.recv_exit_status()
            
            # 读取输出流
            stdout_data = stdout.read().decode('utf-8', errors='ignore').strip()
            stderr_data = stderr.read().decode('utf-8', errors='ignore').strip()
            
            # 关闭流
            stdin.close()
            stdout.close()
            stderr.close()
            
            # 记录执行结果
            if exit_status == 0:
                logger.debug(f"命令执行成功 [{host.name}]: {command}")
            else:
                logger.warning(f"命令执行失败 [{host.name}] 退出码{exit_status}: {command}")
                if stderr_data:
                    logger.debug(f"错误输出: {stderr_data}")
            
            return exit_status, stdout_data, stderr_data
            
        except paramiko.SSHException as e:
            error_msg = f"SSH执行错误 [{host.name}|{host.host}:{host.port}|{host.username}]: {str(e)}"
            logger.error(error_msg)
            raise Exception(error_msg)
        except Exception as e:
            error_msg = f"命令执行异常 [{host.name}|{host.host}:{host.port}|{host.username}] {command}: {str(e)}"
            logger.error(error_msg)
            raise Exception(error_msg)
    
    def _remove_connection(self, connection_key: str):
        """
        从连接池中移除指定连接
        
        Args:
            connection_key: {host.name}|{host.host}:{host.port}|{host.username}
        """
        if connection_key in self._connections:
            try:
                self._connections[connection_key].close()
                self._failed_connections.add(connection_key)
            except:
                pass
            del self._connections[connection_key]
            logger.debug(f"移除连接: {connection_key}")
    
    def close_connection(self, host: Host):
        """
        关闭指定主机的连接
        
        Args:
            host: 主机配置对象
        """
        connection_key = f"{host.name}|{host.host}:{host.port}|{host.username}"

        with self._lock:
            self._remove_connection(connection_key)
    
    def close_all(self):
        """
        关闭所有SSH连接
        
        用于程序结束时的资源清理
        """
        with self._lock:
            logger.info(f"正在关闭 {len(self._connections)} 个SSH连接...")

            # 关闭所有shell
            for key, shell in list(self._shells.items()):
                try:
                    shell.close()
                    logger.info(f"已关闭shell: {key}")
                except Exception as e:
                    logger.warning(f"关闭shell失败 {key}: {str(e)}")
            # 关闭所有连接
            for key, client in list(self._connections.items()):
                try:
                    client.close()
                    logger.info(f"已关闭连接: {key}")
                except Exception as e:
                    logger.warning(f"关闭连接失败 {key}: {str(e)}")
            
            # 清空连接池
            self._connections.clear()
            logger.info("所有SSH连接已关闭")
    
    def get_connection_stats(self) -> Dict[str, any]:
        """
        获取连接池统计信息
        
        Returns:
            Dict[str, Any]: 连接统计信息
        """
        with self._lock:
            return {
                'total_created': self._connection_count,
                'active_connections': len(self._connections),
                'failed_connections': len(self._failed_connections),
                'connection_keys': list(self._connections.keys()),
                'shell_count': len(self._shells),
                'shell_keys': list(self._shells.keys()),
            }
    
    def health_check(self) -> Dict[str, bool]:
        """
        检查所有连接的健康状态
        
        Returns:
            Dict[str, bool]: 连接健康状态映射
        """
        health_status = {}
        
        with self._lock:
            for key, client in self._connections.items():
                health_status[key] = self._is_connection_alive(client)
                if not health_status[key]:
                    self._failed_connections.add(key)

            for key, shell in self._shells.items():
                if shell and not shell.closed and shell.active:
                    health_status[key] = True
                else:
                    health_status[key] = False
                    self._failed_connections.add(key)

        
        return health_status
    
    def __del__(self):
        """
        析构函数：确保资源被正确释放
        """
        try:
            self.close_all()
        except:
            pass  # 析构时忽略异常
    
    def __enter__(self):
        """
        上下文管理器入口
        """
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """
        上下文管理器出口：自动清理资源
        """
        self.close_all()